from collections.abc import Iterable, Iterator, Sequence
from typing import TypeVar, overload

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator
from .counted import CountedDataset
from .iterable import IterableDataset, IterableDatasetIterator
from .sequence import SequenceDataset, SequenceDatasetIterator

AnyKey = TypeVar('AnyKey')
AnyValue = TypeVar('AnyValue')
AnySample = TypeVar('AnySample')


@overload
def wrap_dataset(samples: Dataset[AnySample], /) -> Dataset[AnySample]:
    pass


@overload
def wrap_dataset(samples: Sequence[AnySample], /) -> SequenceDataset[AnySample]:
    pass


@overload
def wrap_dataset(samples: Iterable[AnySample], /) -> IterableDataset[AnySample]:
    pass


@overload
def wrap_dataset(samples: Iterable[AnySample], /, *, samples_n: int) -> CountedDataset[AnySample]:
    pass


def wrap_dataset(samples: Iterable[AnySample], /, *, samples_n: int | None = None) -> Dataset[AnySample]:
    if isinstance(samples, Dataset):
        return samples

    try:
        from datasets import Dataset as HFDataset, IterableDataset as HFIterable
        from .huggingface import HuggingFaceDataset
        if isinstance(samples, (HFDataset, HFIterable)):
            return HuggingFaceDataset(samples)
    except (ImportError, AttributeError):
        pass

    if isinstance(samples, Sequence):
        return SequenceDataset(samples)

    samples = IterableDataset(samples)
    if samples_n is not None:
        samples = CountedDataset(samples, samples_n)
    return samples


@overload
def wrap_dataset_iterator(samples: Iterator[AnySample], /) -> DatasetIterator[AnySample]:
    pass


@overload
def wrap_dataset_iterator(samples: Dataset[AnySample], /) -> DatasetIterator[AnySample]:
    pass


@overload
def wrap_dataset_iterator(samples: Sequence[AnySample], /) -> SequenceDatasetIterator[AnySample]:
    pass


@overload
def wrap_dataset_iterator(samples: Iterable[AnySample], /) -> IterableDatasetIterator[AnySample]:
    pass


def wrap_dataset_iterator(samples: Iterable[AnySample], /) -> DatasetIterator[AnySample]:
    if isinstance(samples, DatasetIterator):
        return samples
    if isinstance(samples, Iterator):
        return IterableDatasetIterator(samples)
    return iter(wrap_dataset(samples))
